# 首先，安装所需的库：
# pip install fastapi uvicorn openai

from openai import OpenAI
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse
import asyncio

app = FastAPI()
client = OpenAI(
    # This is the default and can be omitted
    api_key='sk-b4091c2ef6894f48815cde750a76203b',
    base_url = 'https://api.deepseek.com'
)


html = """
<!DOCTYPE html>
<html>
    <head>
        <title>Chat with OpenAI</title>
    </head>
    <body>
        <h1>WebSocket Chat</h1>
        <form action="" onsubmit="sendMessage(event)">
            <input type="text" id="messageText" autocomplete="off"/>
            <button>Send</button>
        </form>
        <ul id="messages">
        </ul>
        <script>
            var ws = new WebSocket("ws://localhost:8000/ws");
            ws.onmessage = function(event) {
                var messages = document.getElementById('messages');
                var message = document.createElement('li');
                message.textContent = event.data;
                messages.appendChild(message);
            };
            function sendMessage(event) {
                var input = document.getElementById("messageText");
                ws.send(input.value);
                input.value = '';
                event.preventDefault();
            }
        </script>
    </body>
</html>
"""

@app.get("/")
async def get():
    return HTMLResponse(html)

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    try:
        while True:
            prompt = await websocket.receive_text()
            print(f'prompt:{prompt}')
            await openai_streaming_chat(websocket, prompt)
    except WebSocketDisconnect:
        print("Client disconnected")

async def openai_streaming_chat(websocket: WebSocket, prompt: str):
    response=client.chat.completions.create(
    messages=[
        {
            "role": "user",
            "content":prompt,
        }
    ],
    model="deepseek-chat",
    stream=True
)
    for message in response:
        if 'choices' in message:
            await websocket.send_text(message['choices'][0]['delta']['content'])

